import json
import pandas as pd
import numpy as np
import torch
import cv2
import os
import SimpleITK as sitk
import pydicom as dcm
from dataset import heatmap_to_anno
from process_data import dicom2array
from val import predict
from cut_zhuiti_project import get_zhuiti_category
from models.model import model_unet,model_cpn50,model_simple_pose_res18,model_simple_pose_res101,model_resnet_unet,model_DLinknet34



def get_image(data_path):
    json_file = open('Z:\Download/spark_b/spark_b/testB50.json')
    json_file = json.load(json_file)
    series_list = []
    for name in json_file:
        series_Uid = name['seriesUid']
        series_list.append(series_Uid)

    ## 传入一个文件夹路径。从文件夹中筛选中间帧的数据然后转换成image，用来进行预测。同时返回所在的帧数。
    for path,dir_list,file_list in os.walk(data_path):
        for file_name in dir_list:
            print('-----------+++++++++++++++++++++---------------')
            folder_path = os.path.join(path,file_name)
            ## 在一个case的文件夹中遍历所有文件，进行挑选。
            t2_dict = []
            for name in os.listdir(folder_path):
                data_dir = os.path.join(folder_path,name)
                # print(data_dir)
                data = dcm.read_file(data_dir)
                print(series_list)
                print(data.get('SeriesInstanceUID'))
                if data.get('SeriesInstanceUID') in series_list:
                    t2_dict.append(data_dir)
            image_dir_list = t2_dict
            ## 先遍历一遍数据，取出最大的instance num。从而得到中间帧
            max_instance_num = 0
            for dir in image_dir_list:
                instance_num = dcm.read_file(dir).get('InstanceNumber')
                if instance_num > max_instance_num:
                    max_instance_num = instance_num
            index = max_instance_num//2+1
            ## 中间帧的图像可能存在多个，先存dir在list中然后分情况进行处理。
            middle_image_dir = []
            for dir in image_dir_list:
                instance_num = dcm.read_file(dir).get('InstanceNumber')
                if instance_num == index:
                    dir = dir.replace('\\','/')
                    middle_image_dir.append(dir)
            try:
                get_dir = middle_image_dir[0]
            except Exception:
                get_dir = image_dir_list[len(image_dir_list)//2]
                get_dir = get_dir.replace('\\','/')
            image = dicom2array(get_dir)
            print(os.path.join('Z:\Download/spark_b/spark_b/image',get_dir.split('/')[-2] + '_' +
                                  get_dir.split('/')[-1].replace('dcm', 'jpg')))
            cv2.imwrite(os.path.join('Z:\Download/spark_b/spark_b/image',get_dir.split('/')[-2] + '_' +
                                  get_dir.split('/')[-1].replace('dcm', 'jpg')),image)
            # cv2.imshow('image',image)
            # cv2.waitKey()



                # data = sitk.ReadImage(data_dir)
                # t2 = data.GetMetaData('0008|103e')
                # print(t2)
                # if (('t2_tse_sag' in t2) or ('T2' in t2 and 'SAG' in t2 and 'FST2' not in t2) or ('T2' in t2 and 'FSE' in t2 and 'SC' in t2 and "SIIR" not in t2) or ('T2 TES' in t2)
                # or ('T2' in t2 and 'S5mm' in t2) or ('T2' in t2 and 'Sag' in t2) or ('T2 TSE' in t2)):
                # if 'T2' in t2 or 't2' in t2:
                #     print(t2)
            #     if not ('T1' in t2 or 'FGR' == t2 or 'T2_STIR' == t2 or )
            #     # if 't2' in t2 or 'T2' in t2:


            #     if (('t2' in t2 and 'sag' in t2) or ('T2' in t2 and 'SAG' in t2 and 'FSE' not in t2)) and (float(orien[1]) > 0.9 and float(orien[-1]) < -0.9):
            #         data_list.append(data_dir)
            #     elif ('t2' in t2 or 'T2' in t2 ) and (float(orien[1]) > 0.9 and float(orien[-1]) < -0.9):
            #         data_list.append(data_dir)
            #     # elif ('TSE' in t2) and (float(orien[1]) > 0.9 and float(orien[-1]) < -0.9):
            #     #     data_list.append(data_dir)
            #     elif ('FSE' in t2 and 'T1' not in t2) and (float(orien[1]) > 0.9 and float(orien[-1]) < -0.9):
            #         data_list.append(data_dir)

            # ## 取中间位置的数据
            # print(folder_path)

            # print(len(data_list))
            # if len(data_list) < 1 or len(data_list) > 15:
            #     print(folder_path)
            # image_dir = (t2_dict[use_t2[0]])[len(data_list)//2]
            # if len(data_list)> 15:
            #     image_dir = data_list[len(data_list) // 4]
            # image = dicom2array(image_dir)
            # cv2.imshow('image',image)
            # cv2.waitKey()


            # image_dir = data_list[len(data_list)//2]
            # Zindex = len(data_list)//2 + 1

                    # print(t2,data.get('InstanceNumber'))
                    # cv2.imshow('image',image)
                    # cv2.waitKey()

def conver_test_image(test_json_dir,save_path =None):
    # 根据B榜提供的json。直接读取中间帧数据，保存成jpg文件。
    json_file = open(test_json_dir)
    json_file = json.load(json_file)
    series_list = []
    for name in json_file:
        study_uid = name['studyUid']
        series_Uid = name['seriesUid']
        series_list.append(series_Uid)
    data_path = 'Z:\Download/spark_b/spark_b/lumbar_testB50'
    print(series_list)
    for path in os.listdir(data_path):
        for dir in os.listdir(data_path+'/'+path):
            dcm_dir = data_path+'/'+path+'/'+dir
            dcm_data = dcm.read_file(dcm_dir)
            series_Uid = dcm_data.get('SeriesInstanceUID')
            print(series_Uid)
            if series_Uid in series_list:
                image = dicom2array(dcm_dir)
                cv2.imwrite(os.path.join(save_path,str(series_Uid)+'.jpg'),image)







def save_show_test(model,image_size):
    image_path = 'Z:\Download\spark_b\spark_b\image'
    for name in os.listdir(image_path):
        image_dir = os.path.join(image_path, name)
        point_dict = predict(model, image_dir, size=image_size, show_save='../data/show/test_b')



if __name__ == '__main__':

    # get_image(data_path='Z:\Download/spark_b/spark_b/lumbar_testB50')

    # for name in os.listdir('../data/test/study297'):
    #     image_dir = os.path.join('../data/test/study297',name)
    #     data = dcm.read_file(image_dir)
    #     orien = data.get('ImageOrientationPatient')
    #     if orien[1] > 0.9 and orien[-1] < 0.9:
    #         print((data.get('SeriesDescription')).lower())
    #         if 't2' in (data.get('SeriesDescription')).lower():
    #             print('+++++++++++++++++++++')



    show_test_path = 'Z:\Download\spark_b\spark_b\image'
    # Image_Size = (512, 512)
    Image_Size = (768, 768)
    # Image_Size = (1024, 1024)
    # trained_model_dir = '../model_save/Unet_res34_aug/best_.pth'
    # trained_model_dir = '../model_save/DlinkNet34/best_.pth'
    # trained_model_dir = '../model_save/DlinkNet34_768/best_.pth'
    # trained_model_dir = '../model_save/DlinkNet34_1024/best_.pth'
    trained_model_dir = '../model_save/new/DlinkNet34_768/new_.pth'
    # trained_model_dir_2 = '../model_save/DlinkNet34_768/best_.pth'

    # image_path = '../data/show/test_image'
    # model = model_cpn50(image_size = Image_Size)
    # model = model_simple_pose_res101()
    # model = model_resnet_unet(layer='resnet34', pre_train=False, num_class=2)

    model = model_DLinknet34(num_class=2, pre_trained=False)
    model_state = torch.load(trained_model_dir)
    model_state = {k: v for k, v in model_state.items() if 'classify' not in k}
    model.load_state_dict(model_state)

    # model_2 = model_DLinknet34(num_class=2, pre_trained=False, double=False, classify=False)
    # model_2_state = torch.load(trained_model_dir_2)
    # model_2_state = {k: v for k, v in model_2_state.items() if 'classify' not in k}
    # model_2.load_state_dict(model_2_state)

    Test_Aug = True

    save_show_test(model,Image_Size)

    image_dir_list = [
        os.path.join('Z:\Download\spark_b\spark_b\lumbar_testB50', i.replace('\\', '/').split('/')[-1].replace('_', '/').replace('jpg', 'dcm')) for i
        in os.listdir(show_test_path)]
    print(image_dir_list)
    json_file = []
    for image_dir in image_dir_list:
        data = dcm.read_file(image_dir)
        instanceUid = data.get('SOPInstanceUID')
        seriesUid = data.get('SeriesInstanceUID')
        studyUid = data.get('StudyInstanceUID')
        zIndex = data.get('InstanceNumber')
        spacing = data.get('PixelSpacing')[0]
        point_dict = predict(model,image_dir,size=Image_Size,test_aug=Test_Aug,model2 = None)
        single_json = {'data':[],'studyUid':studyUid,"version":"v0.1"}
        single_json['data'].append({"instanceUid":instanceUid,"seriesUid":seriesUid,"annotation":[]})
        single_json['data'][0]['annotation'].append({"annotator": 1,'data':{}})
        single_json['data'][0]['annotation'][0]['data'] = {'point':[]}
        for name,coord in point_dict.items():
            if '-' in name:
                single_json['data'][0]['annotation'][0]['data']['point'].append({
                    'coord':coord,
                    'tag':{
                        'disc':'v1',
                        'identification':str(name)
                    },
                    'zIndex':zIndex,
                })
            else:
                disc = 'vertebra'
                single_json['data'][0]['annotation'][0]['data']['point'].append({
                    'coord': coord,
                    'tag': {
                        'vertebra': 'v2',
                        # 'vertebra': get_zhuiti_category(image_dir,coord,spacing),
                        'identification': str(name)
                    },
                    'zIndex': zIndex,
                })

        json_file.append(single_json)
    print(json_file)
    with open('../data/result_b.json','w') as f:
        json.dump(json_file,f)



    # get_image(data_path='../data/test')


    # t2_dict = {}
    # for dir in os.listdir('../data/test/study222'):
    #     image_dir = os.path.join('../data/test/study222',dir)
    #     image = dcm.read_file(image_dir)
    #     # print(image)
    #     t2 = (image.get('SeriesDescription')).lower()
    #     if t2 not in t2_dict.keys():
    #         t2_dict[t2] = [image_dir]
    #     else:
    #         t2_dict[t2].append(image_dir)
    # print(t2_dict.keys())
    # for image_dir in t2_dict['t2_fse_5mm(s)']:
    #     image = dicom2array(image_dir)
    #     cv2.imshow('image',image)
    #     cv2.waitKey()
        # print(image.get('ImageOrientationPatient'),image.get('SeriesDescription'),image.get('InstanceNumber'))
        # data = sitk.ReadImage(image_dir)
        # t2 = data.GetMetaData('0008|103e')
        # print(t2)

        # # if 'T2' in t2 and 'Sag' in t2:
        # #     print(t2)
        # # print(data)
        # image = dicom2array(image_dir)
        # cv2.imshow('image',image)
        # cv2.waitKey()
